{
 "metadata": {
  "nbsphinx": {
    "execute": "never"
  }
},
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dcd6670a",
   "metadata": {},
   "source": [
    "# Enabling Your GPU for a Solver in MusicBox\n",
    "\n",
    "This tutorial will show you to use utilize a GPU for your MusicBox work.</br>\n",
    "However, this tutorial will not cover how to efficiently use a GPU through parallelization; it will simply introduce getting a GPU set up to run your code.</br>\n",
    "<b>Note:</b> This tutorial requires you to have a Linux GPU-ready environment handy, such as a supercomputing node; it will fail otherwise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e934c67",
   "metadata": {},
   "source": [
    "## 1. Creating a GPU Virtual Environment\n",
    "\n",
    "Running code on a GPU requires a different install protocol when setting up a virtual environment.</br>\n",
    "To do so, run these commands in your terminal:\n",
    "\n",
    "```\n",
    "conda create --name music_box_gpu python=3.9\n",
    "conda activate music_box_gpu\n",
    "pip install --upgrade setuptools pip wheel\n",
    "pip install nvidia-pyindex\n",
    "pip install acom_music_box\n",
    "pip install musica[gpu]\n",
    "conda install ipykernel scikit-learn seaborn scipy dask\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd76107",
   "metadata": {},
   "source": [
    "## 2. Importing MusicBox\n",
    "\n",
    "Importing MusicBox is largerly the same, but with an additional <i>is_cuda_available()</i> function to verify that the GPU is running properly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fe92903",
   "metadata": {},
   "outputs": [],
   "source": [
    "from acom_music_box import MusicBox, Conditions\n",
    "import musica.mechanism_configuration as mc\n",
    "import matplotlib.pyplot as plt\n",
    "from musica.cuda import is_cuda_available\n",
    "from musica import SolverType"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27039537",
   "metadata": {},
   "source": [
    "As with creating the music_box environment in the [Basic Workflow Tutorial](1.%20basic_workflow.ipynb), this cell may be slow to run the first time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d82609a",
   "metadata": {},
   "source": [
    "## 3. Running a Basic Solver on GPU\n",
    "\n",
    "This code is a copy of the [Basic Workflow Tutorial](1.%20basic_workflow.ipynb), but with an if statement added outside the main code to verify that it is running on a GPU.</br>\n",
    "If you are seeing \"Error: No GPU Available\" being printed, that means a GPU was not detected; verify that your environment has a GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54dc9c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "if is_cuda_available():\n",
    "    # Create each of the species that will be simulated\n",
    "    X = mc.Species(name=\"X\")\n",
    "    Y = mc.Species(name=\"Y\")\n",
    "    Z = mc.Species(name=\"Z\")\n",
    "    species = {\"X\": X, \"Y\": Y, \"Z\": Z}\n",
    "    gas = mc.Phase(name=\"gas\", species=list(species.values()))\n",
    "    # Create the reactions that the species undergo in the\n",
    "    arr1 = mc.Arrhenius(name=\"X->Y\", A=4.0e-3, C=50, reactants=[species[\"X\"]], products=[species[\"Y\"]], gas_phase=gas)\n",
    "    arr2 = mc.Arrhenius(name=\"Y->Z\", A=4.0e-3, C=50, reactants=[species[\"Y\"]], products=[species[\"Z\"]], gas_phase=gas)\n",
    "    rxns = {\"X->Y\": arr1, \"Y->Z\": arr2}\n",
    "    # Create the mechanism that is defined by the species, phases, and reactions\n",
    "    mechanism = mc.Mechanism(name=\"tutorial_mechanism\", species=list(species.values()), phases=[gas], reactions=list(rxns.values()))\n",
    "    # Create the box model that contains the mechanism\n",
    "    box_model = MusicBox()\n",
    "    box_model.load_mechanism(mechanism, solver_type=SolverType.cuda_rosenbrock)\n",
    "    # Set the conditions of the box model at time = 0 s\n",
    "    box_model.initial_conditions = Conditions(\n",
    "        temperature=298.15, # Units: Kelvin (K)\n",
    "        pressure=101325.0, # Units: Pascals (Pa)\n",
    "        species_concentrations={ # Units: mol/m^3\n",
    "            \"X\": 3.75,\n",
    "            \"Y\": 5.0,\n",
    "            \"Z\": 2.5,\n",
    "        }\n",
    "    )\n",
    "    # Set the box model conditions at the defined time\n",
    "    box_model.add_evolving_condition(\n",
    "        100.0, # Units: Seconds (s)\n",
    "        Conditions(\n",
    "            temperature=75.0, # Units: Kelvin (K)\n",
    "            pressure=100100.0 # Units: Pascals (Pa)\n",
    "        )\n",
    "    )\n",
    "    # Set the additional configuration options for the box model\n",
    "    box_model.box_model_options.simulation_length = 200 # Units: Seconds (s)\n",
    "    box_model.box_model_options.chem_step_time = 1 # Units: Seconds (s)\n",
    "    box_model.box_model_options.output_step_time = 20 # Units: Seconds (s)\n",
    "    df = box_model.solve()\n",
    "    display(df)\n",
    "    df.plot(x='time.s', y=['CONC.X.mol m-3', 'CONC.Y.mol m-3', 'CONC.Z.mol m-3'], title='Concentration over time', ylabel='Concentration (mol m-3)', xlabel='Time (s)')\n",
    "    plt.show()\n",
    "else:\n",
    "    print(\"Error: No GPU Available\") "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}